import numpy
import torch
import torch.nn as nn
from torch import FloatTensor
import torchvision
from matplotlib import pyplot as plot


class FCN(nn.Module):
    def __init__(self):
        super(FCN, self).__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(5, 5), stride=1, padding=2),
            nn.MaxPool2d(kernel_size=5, stride=2, padding=2),
            nn.LeakyReLU(),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(5, 5), stride=1, padding=2),
            nn.AvgPool2d(kernel_size=(5, 5), stride=2, padding=2),
            nn.LeakyReLU()
        )

        # 2 * 64 * 7 * 7

        self.ccl = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=512, kernel_size=(7, 7), padding=0),  # 512 * 1 * 1
            nn.LeakyReLU(),
            nn.Conv2d(in_channels=512, out_channels=256, kernel_size=(1, 1)),
            nn.LeakyReLU(),
            nn.Conv2d(in_channels=256, out_channels=10, kernel_size=(1, 1)),
            nn.BatchNorm2d(10),
            nn.Sigmoid()
        )

    def forward(self, x):
        cnn_out = self.cnn(x)
        ccl_out = self.ccl(cnn_out)
        y = ccl_out.view(ccl_out.size(0), -1)
        return y


def main():
    draw_x = []
    draw_y = []

    dataset = torchvision.datasets.mnist.MNIST(root="./mnist", download=True)
    fcn = FCN()
    print(fcn)
    loss_function = nn.MSELoss()
    opt = torch.optim.Adam(fcn.parameters(), lr=1e-2)
    train_data = dataset.train_data.resize(60000, 1, 28, 28)
    train_label = dataset.train_labels

    size = 60000
    batch = 5
    page_size = int(size / batch)

    for epoch in range(100):
        for pos in range(page_size):
            start = pos * batch
            end = (pos + 1) * batch

            x = FloatTensor(train_data.numpy()[start:end, :, :, :])
            y = numpy.zeros(shape=(batch, 10))
            yr = train_label[start:end]
            y[numpy.arange(batch), yr] = 1

            y_ = fcn.forward(x)
            loss = loss_function(y_, FloatTensor(y))
            print("LOSS: " + str(loss.detach().numpy()))
            draw_y.append(loss.detach().numpy())
            draw_x.append(epoch * page_size + pos)

            opt.zero_grad()
            loss.backward()
            opt.step()


if __name__ == "__main__":
    main()
